#ifndef SEGMENTED_POOL_H_INCLUDED
#define SEGMENTED_POOL_H_INCLUDED

/*
* Copyright (c) 2010 Michael Collins
*
* This file is part of TerraFirma.
*
* TerraFirma is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* TerraFirma is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with TerraFirma.  If not, see <http://www.gnu.org/licenses/>.
*/

/**segmented_pool.h
 * The segmented pool is an attempt at a statically-sized allocation system like the
 * static_allocator, but better in that it has less memory overhead and pointers to allocated objects
 * are never invalidated by a resize. Hopefully faster too. Should be amortized O(1)
 * Free blocks are initialized to reduce fragmentation instead of LRU
 *
 * Basically, how this works is as a linked-list of arrays, each segment is an array.
 * Random access is achieved by having a vector of pointers to each segment (i.e. a 1-deep tree)
 *
 * Warning -- Only one thread should ever be able to allocate on a certain pool. This isn't thread safe
 * but reads are ok.
 */

#include <vector>
#include <stdint.h>
#include <stdexcept>
#include <stack>

#define CEIL_64(n) ((n) + ((((n)^0x3F)+1) & 0x3F))
#define SEG_IDX(n) ((n) >> 6)
#define SUBSEG_IDX(n) ((n) & 0x3F)

static inline int log_2(uint64_t n)
{
    int res = 0;
    //Hopefully this loop will unroll nicely
    for (char dig = 32; dig > 0; dig >>= 1)
        if (n >> dig)
        {
            n >>= dig;
            res += dig;
        }
    if (n)
        return res+1; //Get last bit
    else
        return res;
}


namespace engine
{
    template<typename T>
    class segmented_pool
    {
    private:
    //public:
        struct segment
        {
            segment() : freebits(0), size(0) {}

            friend class segmented_pool::const_iterator;

            //Destructor is supposed to delete all the elements first
            ~segment()
            {
                for (unsigned int i = 0; i < 64; i++)
                {
                    //Call destructors on all initialized objects
                    if (initialized(i))
                        reinterpret_cast<T*>(&(elements[0][i]))->~T();
                }
            }

            bool initialized(unsigned int n) const
            {
                return (freebits & (1 << n));
            }

            bool empty() const
            {
                return (size==0);
            }

            unsigned int first_element() const
            {
                return log_2(freebits ^ (freebits-1)) - 1;
            }

            unsigned int last_element() const
            {
                return log_2(freebits) - 1;
            }

            //Undefined results if this segment is totally filled up
            const void* next_element_ptr() const //Does not return a T* because the element shouldn't be initialized yet
            {
                //Determine the digit of the least significant unset bit
                int place = log_2(freebits ^ (freebits+1)) - 1;
                return (void*) &elements[0][place];
            }

            //Used for new() operator, allocates and marks the memory as initialized
            void* allocate()
            {
                //Determine the digit of the least significant unset bit
                int place = log_2(freebits ^ (freebits+1)) - 1;
                freebits |= 1 << place; //Set the initialized flag
                size++;
                return (void*) &elements[0][place];
            }

            //Same as above, but can store the subsegment index of the element into arg idx
            void* allocate(unsigned int& idx)
            {
                //Determine the digit of the least significant unset bit
                int place = log_2(freebits ^ (freebits+1)) - 1;
                freebits |= 1 << place; //Set the initialized flag
                size++;
                idx = place;
                return (void*) &elements[0][place];
            }

            //Allocates an object at a certain index, usually used if the object was created on the server
            //and we want our collection to match the servers'
            void* allocate_at(unsigned int idx)
            {
                //If this overwrites an existing object, make sure to delete it
                if (freebits & (1<<idx))
                    reinterpret_cast<T*>(&(elements[0][idx]))->~T(); //Explicitly call destructor
                else
                    size++;

                return (void*) &elements[0][idx];
            }


            //Takes an index argument to an element you want to delete (free)
            void free_idx(unsigned int n)
            {
                if (freebits & (1<<n))
                {
                    //To delete, we need to interpret the element as it's actual type
                    //and call the destructor. this is mildly unsafe, but unless you're
                    //trying to break the system, it should be fine
                    reinterpret_cast<T*>(&(elements[0][n]))->~T(); //Explicitly call destructor
                    freebits ^= (1<<n); //Clear initialized bit, making it free
                    size--;
                }
            }

            //Accessors (STL-like)
            //WARNING -- these will return values regardless of if they are initialized** use with caution
            T& at(unsigned int n)
            {
                if (n > 64) throw std::out_of_range("engine::segmented_pool<T>::segment Range Check");

                return *reinterpret_cast<T*>(&(elements[0][n]));
            }

            const T& at(unsigned int n) const
            {
                if (n > 64) throw std::out_of_range("engine::segmented_pool<T>::segment Range Check");

                return *reinterpret_cast<const T*>(&(elements[0][n]));
            }

            T& operator[](unsigned int n)
            {
                return *reinterpret_cast<T*>(&(elements[0][n]));
            }

            const T& operator[](unsigned int n) const
            {
                return *reinterpret_cast<const T*>(&(elements[0][n]));
            }

            //All pool segments are 64 elements in size
            uint64_t freebits; //A bit vector describing which elements are initialized (1) or not (0)

            uint8_t elements [sizeof(T)][64]; //Uninitialized memory for the elements
            uint8_t size; //Number of initialized spots
        }; //segment

        //segment* test;

        void add_segment()
        {
            segments.push_back(new segment());
            capacity += 64;
        }

        //Adds n segments to the segment array
        void add_segments(unsigned int n)
        {
            segments.reserve(segments.size()+n);
            for (unsigned int i = 0; i < n; i++)
                segments.push_back(new segment());
            capacity += 64*n;
        }

        //Allocates at the location of a previously deleted object
        //undefined results if no object has been deleted
        void* reallocate()
        {
            unsigned int idx = free_idx.top();
            free_idx.pop();
            return segments[SEG_IDX(idx)]->allocate_at(SUBSEG_IDX(idx));
        }

        void* reallocate(unsigned int& idx)
        {
            idx = free_idx.top();
            free_idx.pop();
            return segments[SEG_IDX(idx)]->allocate_at(SUBSEG_IDX(idx));
        }

        //Allocates at the end of a full list
        void* allocate_end()
        {
            if (size==capacity)
                add_segment(); //need room to grow

            unsigned int seg_idx = segments.size() - 1; //We are indexing into the uppermost segment
            return segments[seg_idx]->allocate();
        }

        void* allocate_end(unsigned int& idx)
        {
            if (size==capacity)
                add_segment(); //need room to grow

            unsigned int seg_idx = segments.size() - 1; //We are indexing into the uppermost segment
            unsigned int subseg_idx;
            void* retval = segments[seg_idx]->allocate(subseg_idx);
            idx = (seg_idx << 6) | subseg_idx;
            return retval;
        }

        std::vector<segment*> segments;
        std::stack<unsigned int> free_idx; //Freed spots (previously initialized)
        size_t size, capacity;

    public:
        //Mutable iterator
        //A read-only iterator, but special care is taken so that another thread can add elements
        //to the pool and this iterator will still work
        class const_iterator
        {
        public:
            const_iterator() : subsegment_index(0) {}
            const_iterator(const const_iterator& m) : subsegment_index(m.subsegment_index), ptr(m.ptr) {}
            //const_iterator(const std::vector<segment*>::const_iterator& m, unsigned int idx = 0) :
            //        subsegment_index(idx), it(m) {}

            const_iterator& operator=(const const_iterator& m)
            {
                subsegment_index(m.subsegment_index);
                ptr(m.ptr);
                return *this;
            }

            const T& operator*() const //Dereference operator (supposed to act like STL iterators)
            {
                return (*ptr)->operator[](subsegment_index);
            }

            void operator++()
            {
                //DANGER This has the potential to go running off into uninitialized memory!
                do
                {
                    if (0xA0 & ++subsegment_index)
                    {
                        //overflow the segment
                        subsegment_index = 0;
                        ++ptr;
                    }
                }
                while (!((*ptr)->initialized(subsegment_index)));
            }

            void operator++(int) //postfix
            {
                operator++();
            }

            //Comparative operators
            bool operator< (const const_iterator& m) const
            {
                if (ptr == m.ptr)
                    return (subsegment_index < m.subsegment_index);
                else
                    return (ptr < m.ptr);
            }

            bool operator<= (const const_iterator& m) const
            {
                if (ptr == m.ptr)
                    return (subsegment_index <= m.subsegment_index);
                else
                    return (ptr < m.ptr);
            }

            bool operator> (const const_iterator& m) const
            {
                if (ptr == m.ptr)
                    return (subsegment_index > m.subsegment_index);
                else
                    return (ptr > m.ptr);
            }

            bool operator>= (const const_iterator& m) const
            {
                if (ptr == m.ptr)
                    return (subsegment_index >= m.subsegment_index);
                else
                    return (ptr > m.ptr);
            }

            bool operator== (const const_iterator& m) const
            {
                if (ptr == m.ptr)
                    return (subsegment_index == m.subsegment_index);
                else
                    return false;
            }

            bool operator!= (const const_iterator& m) const
            {
                return !operator==(m);
            }


        private:
            uint8_t subsegment_index;
            const segment * const * ptr; //iterates along the segment pointers in the segments array
        };

        segmented_pool() : size(0), capacity(0) {}

        //Constructor that reserves at least n entries
        segmented_pool(size_t n) : size(0), capacity(0)
        {
            //n_seg is the number of segments required
            size_t n_seg = CEIL_64(n) / 64;
            segments.reserve(n_seg);
            for (size_t i = 0; i < n_seg; i++)
                add_segment();
        }

        ~segmented_pool()
        {
            for (unsigned int it = 0; it < segments.size(); it++)
                delete (segments[it]);
        }

        void* allocate()
        {
            if (!free_idx.empty())
                //We can occupy a previously freed spot!
                return reallocate();
            else
                return allocate_end();
        }

        void* allocate(unsigned int& idx)
        {
            if (!free_idx.empty())
                //We can occupy a previously freed spot!
                return reallocate(idx);
            else
                return allocate_end(idx);
        }

        void* allocate_at(unsigned int idx)
        {
            unsigned int seg_idx = SEG_IDX(idx);
            //If there isn't enough space for this, then you need to grow the collection
            if (seg_idx >= segments.size())
                add_segments(seg_idx - segments.size() + 1);

            return segments[seg_idx]->allocate_at(SUBSEG_IDX(idx));
        }

        //Some convenience method functions for allocating and initializing T objects through common constructors
        //Empty constructor T() needs to be declared in public or friendly access
        T& add()
        {
            return *new(allocate()) T;
        }
        T& add(unsigned int& idx)
        {
            return *new(allocate(idx)) T;
        }
        T& add_at(unsigned int idx)
        {
            return *new(allocate_at(idx)) T;
        }

        //Copy constructed method. This needs the copy constructor of T declared within public or friendly access
        T& add(const T& m)
        {
            return *new(allocate()) T(m);
        }
        T& add(const T& m, unsigned int& idx)
        {
            return *new(allocate(idx)) T(m);
        }
        T& add_at(const T& m, unsigned int idx)
        {
            return *new(allocate_at(idx)) T(m);
        }

        void free(unsigned int idx)
        {
            segments.at(SEG_IDX(idx))->free_idx(SUBSEG_IDX(idx));
            free_idx.push(idx); //We just freed this location, so add it to the free stack
        }

        //Silent free does the same as above, but doesn't add to the free stack
        //This is done on clients that always call allocate_at and add_at
        void free_silent(unsigned int idx)
        {
            segments.at(SEG_IDX(idx))->free_idx(SUBSEG_IDX(idx));
        }

        bool empty () const
        {
            return (size==0);
        }

        //Note: These run in O(n) time! Call begin() and end() sparingly, and if possible cache previous results
        const_iterator begin() const
        {
            if (empty())
                return const_iterator(segments.begin());

            //Linearly find the first initialized object
            const segment* const* it(&(segments[0]));
            for (; it<segments.end(); ++it)
                if (!(*it)->empty())
                    return const_iterator(it, (*it)->first_element());
        }

        const_iterator end() const
        {
            if (empty())
                return const_iterator(segments.begin());

            //Linearly find the last initialized object
            const segment* const* it(segments.end());
            for (; it>segments.begin(); --it)
                if (!(*it)->empty())
                    return const_iterator(it, (*it)->last_element());
        }

        //Accessor functions
        bool is_initialized_segment(unsigned int n) const
        {
            unsigned int seg = SEG_IDX(n);
            if(seg >= segments.size())
                return false; //Out of bounds
            return true;
        }

        const T& at(unsigned int n) const
        {
            if(!is_initialized_segment(n))
                throw std::out_of_range("engine::segmented_pool<T>::at(n): out of segment range");

            return segments[SEG_IDX(n)]->at(SUBSEG_IDX(n));
        }
        T& at(unsigned int n)
        {
            if(!is_initialized_segment(n))
                throw std::out_of_range("engine::segmented_pool<T>::at(n): out of segment range");

            return segments[SEG_IDX(n)]->at(SUBSEG_IDX(n));
        }

        const T& operator[](unsigned int n) const
        {
            return segments[SEG_IDX(n)]->operator[](SUBSEG_IDX(n));
        }
        T& operator[](unsigned int n)
        {
            return segments[SEG_IDX(n)]->operator[](SUBSEG_IDX(n));
        }

        unsigned int next_idx() const
        {
            if (!free_idx.empty())
                return free_idx.top();
            else
                return size; //goes on the end
        }


    };
}


#endif // SEGMENTED_POOL_H_INCLUDED
